package com.ai.controller;


import com.ai.service.VectorStoreService;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.cloud.ai.memory.jdbc.MysqlChatMemoryRepository;
import com.alibaba.cloud.ai.memory.redis.RedisChatMemoryRepository;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

/**
 * RAG检索与查询
 */
@RestController
@RequestMapping("/vectorStore/v1")
@Slf4j
public class VectorStoreController {

    private static final String DEFAULT_PROMPT = "你是一个博学的智能聊天助手，请根据用户提问回答！";

    private final ChatClient dashScopeChatClient;

    @Autowired
    private VectorStoreService vectorStoreService;

    private final VectorStore vectorStore;

    public VectorStoreController(JdbcTemplate jdbcTemplate, ChatClient.Builder chatClientBuilder, VectorStore vectorStore){
        // 构造 ChatMemoryRepository 和 ChatMemory(增加记忆聊天功能)

        // MYSQL记忆存储
        ChatMemoryRepository mysqlChatMemoryRepository = MysqlChatMemoryRepository.mysqlBuilder()
                .jdbcTemplate(jdbcTemplate)
                .build();

        // 内存记忆存储等等 ChatMemoryRepository实现
        InMemoryChatMemoryRepository inMemoryChatMemoryRepository = new InMemoryChatMemoryRepository();

        // 基于Redis 记忆存储
        // RedisChatMemoryRepository redisChatMemoryRepository = RedisChatMemoryRepository.builder().poolConfig(redisConfig);

        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(mysqlChatMemoryRepository)
                .build();

        this.vectorStore = vectorStore;

        this.dashScopeChatClient = chatClientBuilder
                .defaultSystem(DEFAULT_PROMPT)
                // 实现 Logger 的 Advisor
                .defaultAdvisors(new SimpleLoggerAdvisor())
                // 注册Advisor
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                // 设置 ChatClient 中 ChatModel 的 Options 参数
                .defaultOptions(
                        DashScopeChatOptions.builder()
                                .withTopP(0.7)
                                .build()
                )
                .build();
    }

    /**
     * 测试本地向量存储 简单测试
     * @return String
     */
    @GetMapping("/simpleVectorStore")
    public String simpleVectorStore(){
        return vectorStoreService.simpleVectorStore() + "";
    }

    /**
     * 将上传的文件处理后（解析、分割）后存入向量数据库
     * @param file 上传的文件
     * @return 操作结果
     */
    @PostMapping("/vectorize")
    public ResponseEntity<String> vectorizeDocument(@RequestParam("file") MultipartFile file) {
        try {
            //Boolean success = vectorStoreService.saveDocumentToVectorStore(file);
            Boolean success = vectorStoreService.customerSaveDocumentToVectorStore(file);
            //多查询扩展  MultiQueryExpander
            // 查询重写  Query Rewrite  RewriteQueryTransformer.builder() 查询明确化：将模糊的问题转换为具体的查询点
            if (success) {
                return ResponseEntity.ok("文件已成功处理并存入向量数据库");
            } else {
                return ResponseEntity.internalServerError()
                        .body("文件处理失败");
            }
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body("文件处理失败: " + e.getMessage());
        }
    }


    /**
     * 测试阿里百炼知识库向量检索
     * @param userInput String
     * @param conversationId 上下文会话ID
     * @return String
     */
    @GetMapping("/chat")
    public String chatWithDocument(@RequestParam String userInput, @RequestParam(defaultValue = "planet-id") String conversationId,
                                   String sourceFile) {
        // 使用 QuestionAnswerAdvisor，它会自动从VectorStore中检索相关文档作为上下文

        // 自动识别用户问题中提到的文档名
        //String sourceFile = extractSourceFileFromQuestion(userInput);

        SearchRequest searchRequest = SearchRequest.builder()
                .query(userInput)
                // 确保上传解析后的文档向量数据中有 source_file 元数据
                .filterExpression("source_file == '" + sourceFile + "'")
                .topK(5)
                .build();

        QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor
                .builder(vectorStore)
                .searchRequest(searchRequest).build();

        //QuestionAnswerAdvisor advisor = new QuestionAnswerAdvisor(
        //        vectorStore,
        //        searchRequest,
        //        null, // 使用默认的 PromptTemplate
        //        null, // 使用默认的 Scheduler
        //        0     // 使用默认的 order
        //);

        return dashScopeChatClient.prompt()
                .user(userInput)
                .advisors(questionAnswerAdvisor) // 注入检索到的上下文
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, conversationId))
                .call()
                .content();
    }


    /**
     * 根据用户问题检索相关文档并生成回答
     * @param question 用户问题
     * @param topK 返回最相似的文档数量
     * @return AI生成的回答
     */
    @GetMapping("/retrieveAndGenerate")
    public String retrieveAndGenerate(@RequestParam String question, @RequestParam(defaultValue = "5") int topK){
        try {
            // 1. 相似性搜索获取相关文档
            List<Document> relevantDocuments = vectorStoreService.similaritySearch(question, topK);

            if (relevantDocuments.isEmpty()) {
                return "未找到与您的问题相关的信息。";
            }

            // 2. 构建上下文内容
            StringBuilder contextBuilder = new StringBuilder();
            for (int i = 0; i < relevantDocuments.size(); i++) {
                Document doc = relevantDocuments.get(i);
                contextBuilder.append("片段 ").append(i + 1).append(":\n")
                        .append(doc.getText())
                        .append("\n来源: ").append(doc.getMetadata().get("source_file"))
                        .append("\n\n");
            }

            String context = contextBuilder.toString();

            // 3. 构建Prompt
            String promptTemplate = """
                你是智能助手，请基于以下文档内容回答用户的问题。如果文档内容不包含问题的相关信息，请说明无法基于提供的文档回答该问题。
                
                文档内容:
                {context}
                
                用户问题: {question}
                
                请提供准确简洁的回答:
                """;

            String prompt = promptTemplate.replace("{context}", context)
                    .replace("{question}", question);

            // 4. 调用大模型生成回答
            return dashScopeChatClient.prompt(prompt).call().content();

        } catch (Exception e) {
            log.error("RAG检索生成回答时发生错误: ", e);
            return "处理问题时发生错误: " + e.getMessage();
        }
    }
}
